// Copyright 2023 Ant Group Co., Ltd.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once

#include "heu/library/algorithms/paillier_gpu/plaintext.h"
#include "heu/library/algorithms/util/he_object.h"
#include "heu/library/algorithms/util/mp_int.h"

namespace heu::lib::algorithms::paillier_g {

class SecretKey : public HeObject<SecretKey> {
 public:
  Plaintext lambda_;  // lambda = lcm(pâ1, qâ1)
  Plaintext mu_;      // Î¼
  Plaintext p_;       // p
  Plaintext q_;       // q
  // Node: The following paramters are generated by p & q
  Plaintext p_square_;                   // p^2
  Plaintext q_square_;                   // q^2
  Plaintext n_square_;                   // n_ * n_
  Plaintext q_square_inv_mul_q_square_;  // (q^2)^{-1} mod p^2 * q^2
  Plaintext phi_p_square_;               // p(p-1)
  Plaintext phi_q_square_;               // q(q-1)

  void Init();
  // base^exp mod n^2, n = p * q
  Plaintext PowModNSquareCrt(const MPInt &base, const Plaintext &exp) const;

  bool operator==(const SecretKey &other) const {
    return p_ == other.p_ && q_ == other.q_ && lambda_ == other.lambda_ &&
           mu_ == other.mu_;
  }

  bool operator!=(const SecretKey &other) const {
    return !this->operator==(other);
  }

  [[nodiscard]] std::string ToString() const override;
};

}  // namespace heu::lib::algorithms::paillier_g

// clang-format off
namespace msgpack {
MSGPACK_API_VERSION_NAMESPACE(MSGPACK_DEFAULT_API_NS) {
namespace adaptor {

template<>
struct pack<heu::lib::algorithms::paillier_g::SecretKey> {
  template<typename Stream>
  packer<Stream> &operator()(msgpack::packer<Stream> &o,
      const heu::lib::algorithms::paillier_g::SecretKey &sk) const {
    // packing member variables as an array.
    o.pack_array(4);
    o.pack(sk.lambda_);
    o.pack(sk.mu_);
    o.pack(sk.p_);
    o.pack(sk.q_);
    return o;
  }
};

template<>
struct convert<heu::lib::algorithms::paillier_g::SecretKey> {
  msgpack::object const &operator()(const msgpack::object &object,
      heu::lib::algorithms::paillier_g::SecretKey &sk) const {
    if (object.type != msgpack::type::ARRAY) { throw msgpack::type_error(); }
    if (object.via.array.size != 4) { throw msgpack::type_error(); }

    // The order here corresponds to the packer above
    sk.lambda_ = object.via.array.ptr[0].as<heu::lib::algorithms::paillier_g::Plaintext>();
    sk.mu_ = object.via.array.ptr[1].as<heu::lib::algorithms::paillier_g::Plaintext>();
    sk.p_ = object.via.array.ptr[2].as<heu::lib::algorithms::paillier_g::Plaintext>();
    sk.q_ = object.via.array.ptr[3].as<heu::lib::algorithms::paillier_g::Plaintext>();
    sk.Init();
    return object;
  }
};

}  // namespace adaptor
}  // namespace msgpack
}  // namespace msgpack

// clang-format on
